package org.demo3;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.JsonUtils;

import java.util.ArrayList;
import java.util.List;

/**
 * 调用AI服务
 */
public class CallAiService {
    
    private String apiKey = "";

    public String callWithMessage(String question) {
        Generation gen = new Generation();
        List<Message> messages = new ArrayList<>();
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(question).build();
        messages.add(userMsg);
        GenerationParam param =
                GenerationParam.builder().model(Generation.Models.QWEN_TURBO).messages(messages)
                        .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                        .apiKey(apiKey)
                        .build();
        GenerationResult result = null;
        try {
            result = gen.call(param);
        } catch (NoApiKeyException e) {
            throw new RuntimeException(e);
        } catch (InputRequiredException e) {
            throw new RuntimeException(e);
        }
        return result.getOutput().getChoices().get(0).getMessage().getContent();
    }

    /**
     * 姜子牙模型调用
     * @throws NoApiKeyException
     * @throws ApiException
     * @throws InputRequiredException
     */
    public String usage(String question) {
        Generation gen = new Generation();
        GenerationParam param = GenerationParam
                .builder()
                .model("ziya-llama-13b-v1")
                .prompt("<human>:sql改为适配postgresql的语法,只需要返回sql : " + question + "\n<bot>:")
                .build();
        GenerationResult result = null;
        try {
            result = gen.call(param);
        } catch (NoApiKeyException e) {
            throw new RuntimeException(e);
        } catch (InputRequiredException e) {
            throw new RuntimeException(e);
        }
        return JsonUtils.toJson(result);
    }

}
